"""LIBERO环境评估脚本，用于测试训练好的策略在LIBERO任务上的性能。

该脚本会连接到策略服务器，加载LIBERO任务，并执行评估以测试策略性能。
"""

import collections
import dataclasses
import logging
import math
import pathlib

import imageio
from libero.libero import benchmark
from libero.libero import get_libero_path
from libero.libero.envs import OffScreenRenderEnv
import numpy as np
from openpi_client import image_tools
from openpi_client import websocket_client_policy as _websocket_client_policy
import tqdm
import tyro

# LIBERO环境使用的默认动作，包含6个关节动作和1个夹爪动作
LIBERO_DUMMY_ACTION = [0.0] * 6 + [-1.0]
# LIBERO环境渲染使用的分辨率
LIBERO_ENV_RESOLUTION = 256  # resolution used to render training data


@dataclasses.dataclass
class Args:
    """命令行参数配置类，定义了评估脚本所需的各种参数。"""
    
    #################################################################################################################
    # 模型服务器参数
    #################################################################################################################
    # 策略服务器主机地址
    host: str = "0.0.0.0"
    # 策略服务器端口
    port: int = 8000
    # 图像调整大小后的尺寸
    resize_size: int = 224
    # 重新规划动作的步数间隔
    replan_steps: int = 5

    #################################################################################################################
    # LIBERO环境特定参数
    #################################################################################################################
    # 任务套件名称，可选项包括：libero_spatial, libero_object, libero_goal, libero_10, libero_90
    task_suite_name: str = (
        "libero_spatial"  # Task suite. Options: libero_spatial, libero_object, libero_goal, libero_10, libero_90
    )
    # 等待对象在模拟中稳定的步数
    num_steps_wait: int = 10  # Number of steps to wait for objects to stabilize i n sim
    # 每个任务的试验次数
    num_trials_per_task: int = 50  # Number of rollouts per task

    #################################################################################################################
    # 工具参数
    #################################################################################################################
    # 保存视频的输出路径
    video_out_path: str = "data/libero/videos"  # Path to save videos

    # 随机种子（用于可重现性）
    seed: int = 7  # Random Seed (for reproducibility)


def eval_libero(args: Args) -> None:
    """在LIBERO环境中评估策略的主要函数。
    
    Args:
        args: 包含评估配置的参数对象
    """
    # 设置随机种子以确保结果可重现
    np.random.seed(args.seed)

    # 初始化LIBERO任务套件
    benchmark_dict = benchmark.get_benchmark_dict()
    task_suite = benchmark_dict[args.task_suite_name]()
    num_tasks_in_suite = task_suite.n_tasks
    logging.info(f"Task suite: {args.task_suite_name}")

    # 创建视频输出目录
    pathlib.Path(args.video_out_path).mkdir(parents=True, exist_ok=True)

    # 根据任务套件名称设置最大执行步数
    if args.task_suite_name == "libero_spatial":
        max_steps = 220  # 最长的训练演示有193步
    elif args.task_suite_name == "libero_object":
        max_steps = 280  # 最长的训练演示有254步
    elif args.task_suite_name == "libero_goal":
        max_steps = 300  # 最长的训练演示有270步
    elif args.task_suite_name == "libero_10":
        max_steps = 520  # 最长的训练演示有505步
    elif args.task_suite_name == "libero_90":
        max_steps = 400  # 最长的训练演示有373步
    else:
        raise ValueError(f"Unknown task suite: {args.task_suite_name}")

    # 创建WebSocket客户端连接到策略服务器
    client = _websocket_client_policy.WebsocketClientPolicy(args.host, args.port)

    # 开始评估
    total_episodes, total_successes = 0, 0
    # 遍历任务套件中的每个任务
    for task_id in tqdm.tqdm(range(num_tasks_in_suite)):
        # 获取当前任务
        task = task_suite.get_task(task_id)

        # 获取LIBERO默认的初始状态
        initial_states = task_suite.get_task_init_states(task_id)

        # 初始化LIBERO环境和任务描述
        env, task_description = _get_libero_env(task, LIBERO_ENV_RESOLUTION, args.seed)

        # 开始执行试验
        task_episodes, task_successes = 0, 0
        # 对每个任务执行多次试验
        for episode_idx in tqdm.tqdm(range(args.num_trials_per_task)):
            logging.info(f"\nTask: {task_description}")

            # 重置环境
            env.reset()
            # 动作计划队列，用于存储动作块
            action_plan = collections.deque()

            # 设置初始状态
            obs = env.set_init_state(initial_states[episode_idx])

            # 初始化变量
            t = 0
            # 用于保存回放视频的图像列表
            replay_images = []

            logging.info(f"Starting episode {task_episodes+1}...")
            # 执行episode直到达到最大步数
            while t < max_steps + args.num_steps_wait:
                try:
                    # 重要：在前几个时间步什么都不做，因为模拟器会掉落对象
                    # 我们需要等待它们稳定下来
                    if t < args.num_steps_wait:
                        obs, reward, done, info = env.step(LIBERO_DUMMY_ACTION)
                        t += 1
                        continue

                    # 获取预处理后的图像
                    # 重要：旋转180度以匹配训练时的预处理
                    img = np.ascontiguousarray(obs["agentview_image"][::-1, ::-1])
                    wrist_img = np.ascontiguousarray(obs["robot0_eye_in_hand_image"][::-1, ::-1])
                    img = image_tools.convert_to_uint8(
                        image_tools.resize_with_pad(img, args.resize_size, args.resize_size)
                    )
                    wrist_img = image_tools.convert_to_uint8(
                        image_tools.resize_with_pad(wrist_img, args.resize_size, args.resize_size)
                    )

                    # 保存预处理后的图像用于回放视频
                    replay_images.append(img)

                    # 如果动作计划为空，则计算新的动作块
                    if not action_plan:
                        # 完成了之前动作块的执行--计算新的动作块
                        # 准备观测字典
                        element = {
                            "observation/image": img,
                            "observation/wrist_image": wrist_img,
                            "observation/state": np.concatenate(
                                (
                                    obs["robot0_eef_pos"],
                                    _quat2axisangle(obs["robot0_eef_quat"]),
                                    obs["robot0_gripper_qpos"],
                                )
                            ),
                            "prompt": str(task_description),
                            "viewmats": obs.get("viewmats", None),
                            "Ks": obs.get("Ks", None),
                        }

                        # 查询模型获取动作
                        action_chunk = client.infer(element)["actions"]
                        assert (
                            len(action_chunk) >= args.replan_steps
                        ), f"We want to replan every {args.replan_steps} steps, but policy only predicts {len(action_chunk)} steps."
                        action_plan.extend(action_chunk[: args.replan_steps])

                    # 从动作计划中取出下一个动作
                    action = action_plan.popleft()

                    # 在环境中执行动作
                    obs, reward, done, info = env.step(action.tolist())
                    if done:
                        task_successes += 1
                        total_successes += 1
                        break
                    t += 1

                except Exception as e:
                    logging.error(f"Caught exception: {e}")
                    break

            task_episodes += 1
            total_episodes += 1

            # 保存episode的回放视频
            suffix = "success" if done else "failure"
            task_segment = task_description.replace(" ", "_")
            imageio.mimwrite(
                pathlib.Path(args.video_out_path) / f"rollout_{task_segment}_{suffix}.mp4",
                [np.asarray(x) for x in replay_images],
                fps=10,
            )

            # 记录当前结果
            logging.info(f"Success: {done}")
            logging.info(f"# episodes completed so far: {total_episodes}")
            logging.info(f"# successes: {total_successes} ({total_successes / total_episodes * 100:.1f}%)")

        # 记录当前任务的最终结果
        logging.info(f"Current task success rate: {float(task_successes) / float(task_episodes)}")
        logging.info(f"Current total success rate: {float(total_successes) / float(total_episodes)}")

    # 记录总体结果
    logging.info(f"Total success rate: {float(total_successes) / float(total_episodes)}")
    logging.info(f"Total episodes: {total_episodes}")


def _get_libero_env(task, resolution, seed):
    """初始化并返回LIBERO环境，以及任务描述。
    
    Args:
        task: LIBERO任务对象
        resolution: 环境渲染分辨率
        seed: 随机种子
        
    Returns:
        tuple: (环境对象, 任务描述)
    """
    task_description = task.language
    task_bddl_file = pathlib.Path(get_libero_path("bddl_files")) / task.problem_folder / task.bddl_file
    env_args = {"bddl_file_name": task_bddl_file, "camera_heights": resolution, "camera_widths": resolution}
    env = OffScreenRenderEnv(**env_args)
    env.seed(seed)  # 重要：种子似乎会影响即使使用固定初始状态时对象的位置
    
    # 获取相机参数
    def env_step_with_camera_params(action):
        obs, reward, done, info = env.step(action)
        # 添加相机参数
        if "agentview_image" in obs:
            # 获取相机视图矩阵和内参矩阵
            camera_id = env.sim.model.camera_name2id("agentview")
            viewmat = env.sim.model.camera_xpos[camera_id].copy()
            K = np.eye(3)  # 简化的相机内参矩阵
            K[0, 0] = K[1, 1] = env.sim.model.cam_fovy[camera_id]
            obs["viewmats"] = np.array([viewmat])
            obs["Ks"] = np.array([K])
        return obs, reward, done, info
    
    # 替换环境的step方法
    env.original_step = env.step
    env.step = env_step_with_camera_params
    


def _quat2axisangle(quat):
    """将四元数转换为轴角表示。
    
    从robosuite复制：https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55
    
    Args:
        quat: 四元数[x, y, z, w]
        
    Returns:
        轴角表示的旋转
    """
    # 裁剪四元数
    if quat[3] > 1.0:
        quat[3] = 1.0
    elif quat[3] < -1.0:
        quat[3] = -1.0

    den = np.sqrt(1.0 - quat[3] * quat[3])
    if math.isclose(den, 0.0):
        # 这是一个（接近）零度旋转，立即返回
        return np.zeros(3)

    return (quat[:3] * 2.0 * math.acos(quat[3])) / den


if __name__ == "__main__":
    # 配置日志级别
    logging.basicConfig(level=logging.INFO)
    # 使用tyro解析命令行参数并运行评估函数
    tyro.cli(eval_libero)